# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import collections
import logging

import tabulate

from coin.proto.coin_query_pb2 import ProductFillElement


class FillChecker(object):
  def __init__(self, og_info, market_type, *, num_allowed_fails=12, rel_tol=1e-2, abs_tol=1e-4):
    assert market_type in ('Spot', 'Futures', 'Margin'), market_type
    self._og_info = og_info
    self._market_type = market_type
    self._rel_tol = rel_tol
    self._abs_tol = abs_tol
    self._fill_og_info = og_info.clone()
    self._check_stat = collections.deque(maxlen=num_allowed_fails + 1)
    self._init_done = False
    self._logger = logging.getLogger(__name__)

  def check(self):
    is_fail = not self._fill_og_info.isclose(
        self._og_info, rel_tol=self._rel_tol, abs_tol=self._abs_tol)
    self._check_stat.append(is_fail)
    return not all(self._check_stat)

  def _copy_balance(self):
    for currency in self._og_info.currencies:
      balance = self._og_info.get_balance(currency)
      self._fill_og_info.set_balance(currency, balance)

  def _copy_position(self):
    for product in self._og_info.products:
      position = self._og_info.get_position(product)
      self._fill_og_info.set_position(product, position)

  def reinit(self):
    self._check_stat.clear()
    if self._market_type == 'Spot':
      self._copy_balance()
    elif self._market_type in ('Futures', 'Margin'):
      self._copy_position()
    else:
      raise ValueError('Unknown market type! %s', self._market_type)
    self._init_done = True

  def is_initialized(self):
    return self._init_done

  def comparison_table_as_dict(self):
    if self._market_type == 'Spot':
      return self._spot_check_table_as_dict()
    elif self._market_type in ('Futures', 'Margin'):
      return self._position_check_table_as_dict()
    else:
      raise ValueError('Unknown market type: %s', self._market_type)

  def _spot_check_table_as_dict(self):
    balance = collections.OrderedDict()
    fill_balance = collections.OrderedDict()
    balance['name'] = 'Balance'
    fill_balance['name'] = 'Balance(fill)'

    for currency in self._og_info.currencies:
      currency_str = currency.currency
      balance[currency_str] = self._og_info.get_balance(currency)
      fill_balance[currency_str] = self._fill_og_info.get_balance(currency)
    return [balance, fill_balance]

  def _position_check_table_as_dict(self):
    position = collections.OrderedDict()
    fill_position = collections.OrderedDict()
    position['name'] = 'Position'
    fill_position['name'] = 'Position(fill)'
    for product in self._og_info.products:
      symbol = product.symbol
      position[symbol] = self._og_info.get_position(product)
      fill_position[symbol] = self._fill_og_info.get_position(product)
    return [position, fill_position]

  def get_formatted_comparison_table(self):
    table = tabulate.tabulate(self.comparison_table_as_dict(), headers='keys')
    return 'Fill check table (%s-%s):\n %s' % (self._market_type, self._og_info.exchange, table)

  def update_by_fill(self, fill_proto: ProductFillElement):
    self._fill_og_info.update_by_fill(fill_proto)
